import torch
from torch import nn
from torchvision import transforms
from torchvision.datasets import ImageFolder
from NumberNet import NumberNet
import matplotlib.pyplot as plt

root_dir = "./data/numbers"
transform = transforms.Compose([
    transforms.Resize(30),  # 将原本的图像进行改变大小
    transforms.CenterCrop(28),  # 裁剪图像区域中心位置的内容
    transforms.RandomRotation(15),  # 让图像随机旋转0~15°
    transforms.RandomHorizontalFlip(),  # 让图像进行随机的水平翻转
    transforms.ToTensor()
])
number_datasets = ImageFolder(root_dir, transform=transform)
print(number_datasets)
num = len(number_datasets)
a = int(num * 0.8)  # train 数据集的长度
b = num - a  # valid 数据集的长度
# 将原本的图像数据集拆分成两份，train_datasets、valid_datasets
train_datasets, valid_datasets = torch.utils.data.random_split(number_datasets, [a, b])

train_loader = torch.utils.data.DataLoader(train_datasets, batch_size=1000, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_datasets, batch_size=b)

# -----------------模型设计-------------------
model = NumberNet(classes=10)
# -----------------损失函数-------------------
criterion = nn.CrossEntropyLoss()  # 交叉熵 + softmax
# -----------------优化器-------------------
sgd = torch.optim.SGD(model.parameters(), lr=0.1)
# -----------------训练图像-----------------
epochs = 10
for epoch in range(epochs):
    print(f"----epoch {epoch + 1} / {epochs}:")
    for images, labels in iter(train_loader):
        sgd.zero_grad()
        predict_labels = model(images.to(model.device))
        # -----------对模型的损失率进行计算---------------
        loss = criterion(predict_labels, labels.to(model.device))
        loss.backward()
        sgd.step()
        # -----------对模型的准确率进行计算---------------
        with torch.no_grad():
            valid_images, valid_labels = next(iter(valid_loader))
            valid_predict_labels = model(valid_images.to(model.device))
            valid_predict = torch.argmax(valid_predict_labels, dim=-1)
            accuracy = sum(valid_predict == valid_labels.to(model.device)) / valid_labels.shape[0]
        print(f"------loss:{loss.item():.4f} ------- accuracy:{accuracy.item():.4f}")

torch.save(model, "numbers_model2.pth")
